# coding: utf-8
import os
import numpy
import numpy as np
import matplotlib.pyplot as plt
from pandas import read_csv
import math
import pandas as pd
from keras.models import Sequential
from keras.layers import Dense
from keras.callbacks import EarlyStopping
from keras.layers import LSTM
from keras.layers import RNN, SimpleRNN, SimpleRNNCell
from sklearn.preprocessing import MinMaxScaler
from sklearn.metrics import mean_squared_error, mean_absolute_error, r2_score
from keras.utils.vis_utils import plot_model


# 创建数据集pi
def create_dataset(dataset, scale):
    train_size = int(len(dataset) * scale)
    # 将数据集分为两部分
    train, test = dataset[0:train_size], dataset[train_size:]
    # train, test = dataset.iloc[0:train_size], dataset.iloc[train_size:]

    return train, test


def set_label_dataset(dataset, scale):
    dataset_cut = int(scale)
    x_dataset, y_dataset = dataset[:-dataset_cut], dataset[dataset_cut:]
    # x_dataset, y_dataset = dataset.iloc[:-dataset_cut], dataset.iloc[dataset_cut:]
    return x_dataset, y_dataset


def mape(y_true, y_pred):
    return numpy.mean(numpy.abs((y_pred - y_true) / y_true)) * 100


def mean_absolute_percentage_error(y_true, y_pred):
    # 平均绝对百分比误差（MAPE）的计算
    y_true, y_pred = np.array(y_true), np.array(y_pred)
    return np.mean(np.abs((y_true - y_pred) / y_true)) * 100


def get_r2_numpy(x, y):
    # polyfit为最小二乘多项式拟合，返回多项式的系数矩阵，和协方差矩阵（？）
    slope, intercept = np.polyfit(x, y, 1)
    r_squared = 1 - (sum((y - (slope * x + intercept)) ** 2) / ((len(y) - 1) * np.var(y, ddof=1)))
    return r_squared


def score_R2(ans, y_test):
    pmean = np.mean(y_test)
    omean = np.mean(ans)
    SSR = 0.0
    varp = 0.0
    varo = 0.0
    for i in range(0, len(y_test)):
        diffXXbar = y_test[i] - pmean
        difYYbar = ans[i] - omean
        SSR += (diffXXbar * difYYbar)
        varo += diffXXbar ** 2
        varp += difYYbar ** 2
    SST = math.sqrt(varo * varp)
    return (SSR / SST) ** 2

if __name__ == '__main__':
    WORK_PATH = os.getcwd()  # 获取当前项目工作目录，
    DATA_PATH = os.path.join(WORK_PATH, r'indices_caculate\result\ROW_SPEI-12\SPEI-12_52533.txt')  # models 目录

    # 加载数据
    dataframe = read_csv(DATA_PATH, header=None, names=('TIME','SPEI-1'))
    dataframe = dataframe.set_index(['TIME'], drop=True) # 把日期作为索引


    # 创建训练和测试数据集
    train, test = create_dataset(np.array(dataframe), 0.9)

    # 将数据归一化处理???为什么需要归一化
    # MaxAbsScaler是归到[-1,1],MinMaxScaler归到[0,1]
    scaler = MinMaxScaler(feature_range=(0, 1))
    scaler.fit(train) # 从train的数据中学到缩放规则
    train = scaler.transform(train)
    test = scaler.transform(test)


    # 获取标签数据，即，根据预测的范围通过平移scale来对应预测的对象，根据x预测y
    X_train, Y_train = set_label_dataset(train, 1)
    X_test, Y_test = set_label_dataset(test, 1)

    """
    1.创建Sequential模型
    2.通过.add（）方法添加层
    3.将input_dim参数传递给第一层
    4.设置激活函数为Relu
    5.设置损失函数为均方误差mean_squared_error，优化器为adam
    6.当监控到损失loss不再减少时停止训练
    7.patience=2表示损失没有改善的训练周期数,在这之后将停止训练
    8.ANN将训练epochs个周期
    9.配置完上述参数后,编译模型
    """
    ANN_model = Sequential()
    """
    1.这里的Dense的第一个参数是节点数，可以从1~无穷调，可以加很多层Dence，这是第一个需要搜索的地方
    """
    ANN_model.add(Dense(12, input_dim=X_train.shape[1],activation='relu'))
    # ANN_model.add(Dense(12, input_dim=X_train.shape[1], activation='relu'))
    ANN_model.add(Dense(1))
    ANN_model.summary()
    """
    2.这里的optimizer:[SGD,RMSprop,Adagrad,Adadelta,Adam,Adamax,Nadam]
    epochs:[500,1000,2000,3000,4000,5000,6000,7000,8000] ,可以把early_stop去掉来试试不同的epochs
    batch_size 可以从1~无穷，主要是看机器能带动的程度
    """
    ANN_model.compile(loss='mean_squared_error', optimizer='RMSprop', metrics=['accuracy'])
    early_stop = EarlyStopping(monitor='loss', patience=5, verbose=1)
    history = ANN_model.fit(X_train, Y_train, epochs=5000, batch_size=500, callbacks=[early_stop], verbose=1, shuffle=False)

    y_train_pred = ANN_model.predict(X_train)
    y_test_pred = ANN_model.predict(X_test)

    y_train_true = scaler.inverse_transform(Y_train)
    y_pred_true = scaler.inverse_transform(Y_test)
    y_train = scaler.inverse_transform(y_train_pred)
    y_pred = scaler.inverse_transform(y_test_pred)

    print("train rmse: ", numpy.sqrt(mean_squared_error(y_train_true, y_train)))
    print("train mae: ", mean_absolute_error(y_train_true, y_train))
    print("train mape: ", mean_absolute_percentage_error(y_train_true, y_train))
    print("train r2: ", r2_score(y_train_true, y_train))

    print("predict rmse: ", numpy.sqrt(mean_squared_error(y_pred_true,y_pred)))
    print("predict mae: ", mean_absolute_error(y_pred_true,y_pred))
    print("predict mape: ", mean_absolute_percentage_error(y_pred_true, y_pred))
    print("predict r2: ", r2_score(y_pred_true, y_pred))

    print('预测的数据个数：',len(y_pred))

    true_list = numpy.array(y_pred_true)
    pred_list = numpy.array(y_pred)

    # squeeze()是去掉维度为1的维度
    # print("true_list: " + str(list(true_list.squeeze())))
    # print("pred_list: " + str(list(pred_list.squeeze())))

    fig = plt.figure()
    ax = fig.add_subplot(1, 1, 1)
    ax.plot(range(len(true_list)), true_list, c="black")
    ax.plot(range(len(pred_list)), pred_list, c="red")
    # x_ticks = ax.set_xticks(range(len(true_list) / 24, len(true_list), len(true_list) / 12))
    # x_labels = ax.set_xticklabels(range(1, 13, 1), fontsize="small")
    ax.set_xlabel("Month")
    ax.set_ylabel("SPEI-1")
    ax.legend(['true', 'pred'])
    plt.show()



